﻿using System;
using System.Collections.Generic;
using System.Linq;
using Algoritmia.DataStructures.Digraphs;
using Algoritmia.DataStructures.Queues;
using Algoritmia.Utils;
using Sets = Algoritmia.DataStructures.Sets;

namespace Algoritmia.Problems.Traversals
{
    public class RecursiveDepthFirstTraverser<TVertex,TResult> : IGraphDepthFirstTraverser<TVertex,TResult>
    {
        private readonly Func<IEnumerable<TVertex>, Sets.ISet<TVertex>> _createSet;
        private readonly Func<IEnumerable<TVertex>, IFifo<TResult>> _createFifo;

        public RecursiveDepthFirstTraverser(Func<IEnumerable<TVertex>,IFifo<TResult>> createFifo = null,
            Func<IEnumerable<TVertex>,Sets.ISet<TVertex>> createSet = null)
        {
            _createSet = createSet ?? (items => AlgoritmiaFactory.Sets<TVertex>.Default);
            _createFifo = createFifo ?? (items => AlgoritmiaFactory.Fifos<TResult>.Default);
        }

        private static void InternalTraverse(IDigraph<TVertex> graph, TVertex u, TVertex v,
                                                      Sets.ISet<TVertex> visited, IFifo<TResult> queue, 
                                                      Func<TVertex, TVertex, TResult> preorderVisitor, 
                                                      Func<TVertex, TVertex, TResult> postorderVisitor)
        {
            visited.Add(v);
            if ( preorderVisitor != null)
            {
                queue.Push(preorderVisitor(u,v));
            }
            foreach (var w in graph.Successors(v))
            {
                if (!visited.Contains(w))
                {
                    InternalTraverse(graph, v, w, visited, queue, preorderVisitor, postorderVisitor);
                }
            }
            if (postorderVisitor != null)
            {
                queue.Push(postorderVisitor(u, v));
            }
        }


        private static Func<TVertex, TVertex, TResult> GetPreorderVisitorFunction(Func<TVertex, TVertex, TResult> preorderVisitor, Func<TVertex, TVertex, TResult> postorderVisitor)
        {
            if (preorderVisitor == null && postorderVisitor == null)
            {
                preorderVisitor = (u, v) => (TResult)(object)v;
            }
            return preorderVisitor;
        }

        #region IGraphDepthFirstTraverser<TVertex,TResult> Members

        public IEnumerable<TResult> Traverse(IDigraph<TVertex> digraph, TVertex source, 
            Func<TVertex, TVertex, TResult> preorderVisitor = null, 
            Func<TVertex, TVertex, TResult> postorderVisitor = null)
        {
            preorderVisitor = GetPreorderVisitorFunction(preorderVisitor, postorderVisitor);
            var visited = _createSet(digraph.V);
            var fifo = _createFifo(digraph.V);
            InternalTraverse(digraph, source, source, visited, fifo, preorderVisitor, postorderVisitor);
            return fifo;
        }

        public IEnumerable<TResult> FullTraverse(IDigraph<TVertex> digraph, 
            Func<TVertex, TVertex, TResult> preorderVisitor, 
            Func<TVertex, TVertex, TResult> postorderVisitor)
        {
            preorderVisitor = GetPreorderVisitorFunction(preorderVisitor, postorderVisitor);
            var visited = _createSet(digraph.V);
            var fifo = _createFifo(digraph.V);
            foreach(var u in digraph.V)
            {
                if (!visited.Contains(u))
                {
                    InternalTraverse(digraph,u,u,visited,fifo,preorderVisitor,postorderVisitor);
                }
                while ( fifo.Count > 0)
                {
                    yield return fifo.Pop();
                }
            }
        }

        #endregion
    }


    public class DepthFirstTraverser<TVertex,TResult> : IGraphDepthFirstTraverser<TVertex,TResult>
    {

        private readonly Func<IEnumerable<TVertex>, Sets.ISet<TVertex>> _createSet;
        private readonly Func<IEnumerable<TVertex>, ILifo<Tuple<TVertex,TVertex>>> _createLifo;

        public DepthFirstTraverser(Func<IEnumerable<TVertex>,Sets.ISet<TVertex>> createSet = null,
                                   Func<IEnumerable<TVertex>,ILifo<Tuple<TVertex,TVertex>>> createFifo = null)
        {
            _createSet = createSet ?? (items => AlgoritmiaFactory.Sets<TVertex>.Default);
            _createLifo = createFifo ?? (items => AlgoritmiaFactory.Lifos<Tuple<TVertex,TVertex>>.Default);
        }

        #region Private Methods

        private static IEnumerable<TResult> InternalTraverse(IDigraph<TVertex> digraph, TVertex source, 
                                                      Sets.ISet<TVertex> visited, ILifo<Tuple<TVertex,TVertex>> stack, 
                                                      Func<TVertex, TVertex, TResult> preorderVisitor, 
                                                      Func<TVertex, TVertex, TResult> postorderVisitor)
        {
            visited.Add(source);
            if (preorderVisitor != null)
            {
                yield return preorderVisitor(source, source);
            }
            stack.Push(new Tuple<TVertex, TVertex>(source, source));
            while (stack.Count > 0)
            {
                var edge = stack.Pop();
               
                TVertex w = default(TVertex);
                
                bool allSuccessorsHaveBeenVisited = false;
                try
                {
                    w = digraph.Successors(edge.Item2).First( v => !visited.Contains(v) );
                }
                catch (InvalidOperationException)
                {
                    allSuccessorsHaveBeenVisited = true;
                }

                if (!allSuccessorsHaveBeenVisited)
                {
                    stack.Push(edge);
                    stack.Push(new Tuple<TVertex,TVertex>(edge.Item2,w));
                    visited.Add(w);
                    if (preorderVisitor != null)
                    {
                        yield return preorderVisitor(edge.Item2, w);
                    }
                }
                else
                {
                    if (postorderVisitor != null)
                        yield return postorderVisitor(edge.Item1, edge.Item2);
                }
            }
        }

        #endregion

        #region IGraphDepthFirstTraverser<TVertex,TResult> Members

        public IEnumerable<TResult> Traverse(IDigraph<TVertex> digraph, TVertex source, 
                                             Func<TVertex, TVertex, TResult> preorderVisitor = null, 
                                             Func<TVertex, TVertex, TResult> postorderVisitor = null)
        {
            if (preorderVisitor == null && postorderVisitor == null)
            {
                preorderVisitor = (v1, v2) => (TResult) (object) v2;
            }
            var visited = _createSet(digraph.V);
            var stack = _createLifo(digraph.V);
            return InternalTraverse(digraph, source, visited,stack, preorderVisitor, postorderVisitor);
        }

        public IEnumerable<TResult> FullTraverse(IDigraph<TVertex> digraph, Func<TVertex, TVertex, TResult> preorderVisitor = null, Func<TVertex, TVertex, TResult> postorderVisitor = null)
        {
            var visited = _createSet(digraph.V);
            var stack = _createLifo(digraph.V);
            return digraph.V.Where(v => !visited.Contains(v)).SelectMany(v => InternalTraverse(digraph, v, visited, stack, preorderVisitor, postorderVisitor));
        }

        #endregion
    }
}
